from ultralytics import YOLO
import torch

# 确保 MPS 可用
if not torch.backends.mps.is_available():
    raise RuntimeError("MPS backend is not available.")

# 设置设备为 MPS
device = 'mps'  

# Load a COCO-pretrained YOLOv8n model
model = YOLO('./yolov8n.pt')

# Train the model on your dataset for 100 epochs
results = model.train(data='./data.yaml', epochs=100, fliplr=0, device=device)  # 添加 device 参数